#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
华为云垃圾分类项目环境设置脚本
用于检查和安装项目所需的依赖包
"""

import sys
import subprocess
import pkg_resources
import platform

def check_python_version():
    """检查Python版本"""
    print("检查Python版本...")
    version = sys.version_info
    if version.major != 3 or version.minor != 6:
        print(f"❌ 错误: 当前Python版本为 {version.major}.{version.minor}")
        print("   本项目需要Python 3.6.x版本")
        print("   请使用以下命令创建Python 3.6环境:")
        print("   conda create -n garbage_classify python=3.6")
        print("   conda activate garbage_classify")
        return False
    else:
        print(f"✅ Python版本正确: {version.major}.{version.minor}.{version.micro}")
        return True

def get_required_packages():
    """获取所需的包列表"""
    return {
        'tensorflow': '1.15.0',
        'keras': '2.3.1',
        'Pillow': '5.0.0',
        'numpy': '1.16.4',
        'scikit-learn': '0.21.3',
        # 'moxing': '1.17.0',  # 可选，仅用于华为云环境
        'opencv-python': '4.1.2.30',
        'matplotlib': '3.1.2',
        'h5py': '2.10.0'
    }

def check_package(package_name, required_version):
    """检查单个包的版本"""
    try:
        installed_version = pkg_resources.get_distribution(package_name).version
        if installed_version == required_version:
            print(f"✅ {package_name} {installed_version}")
            return True
        else:
            print(f"⚠️  {package_name} {installed_version} (需要 {required_version})")
            return False
    except pkg_resources.DistributionNotFound:
        print(f"❌ {package_name} 未安装")
        return False

def install_package(package_name, version):
    """安装指定版本的包"""
    try:
        package_spec = f"{package_name}=={version}"
        print(f"正在安装 {package_spec}...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", package_spec])
        print(f"✅ {package_spec} 安装成功")
        return True
    except subprocess.CalledProcessError as e:
        print(f"❌ {package_spec} 安装失败: {e}")
        return False

def main():
    """主函数"""
    print("=" * 60)
    print("华为云垃圾分类项目环境检查工具")
    print("=" * 60)
    
    # 检查Python版本
    if not check_python_version():
        sys.exit(1)
    
    print("\n检查已安装的包...")
    required_packages = get_required_packages()
    missing_packages = []
    wrong_version_packages = []
    
    for package, version in required_packages.items():
        if not check_package(package, version):
            missing_packages.append((package, version))
    
    if not missing_packages:
        print("\n🎉 所有依赖包都已正确安装！")
        print("可以开始使用项目了。")
        return
    
    print(f"\n发现 {len(missing_packages)} 个包需要安装或更新")
    
    # 询问是否自动安装
    response = input("\n是否自动安装缺失的包？(y/n): ").lower().strip()
    if response not in ['y', 'yes', '是']:
        print("请手动安装缺失的包:")
        for package, version in missing_packages:
            print(f"pip install {package}=={version}")
        return
    
    # 自动安装
    print("\n开始安装缺失的包...")
    failed_packages = []
    
    for package, version in missing_packages:
        if not install_package(package, version):
            failed_packages.append((package, version))
    
    if failed_packages:
        print(f"\n❌ {len(failed_packages)} 个包安装失败:")
        for package, version in failed_packages:
            print(f"   {package}=={version}")
        print("\n请手动安装这些包:")
        for package, version in failed_packages:
            print(f"pip install {package}=={version}")
    else:
        print("\n🎉 所有包安装成功！")
        print("可以开始使用项目了。")

if __name__ == "__main__":
    main() 